import json
import os
from flask import current_app
import torch
from transformers import BertTokenizer, BertConfig


class Config:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.data_path = os.path.join(current_app.root_path,'data')
        self.save_model_path = os.path.join(current_app.root_path,'save_model')

        self.tag2id_path = os.path.join(self.data_path, 'tag2id.json')
        self.label_path = os.path.join(self.data_path, 'labels.json')
        self.train_txt_path = os.path.join(self.data_path, 'train.txt')
        self.vocab_txt_path = os.path.join(self.data_path, 'vocab.txt')
        self.origin_path = os.path.join(current_app.root_path, 'data_origin')
        self.model_name = 'BiLSTM_CRF'
        self.embedding_dim = 300
        self.hidden_dim = 256
        self.dropout =0.2
        self.lr = 2e-3
        self.epochs = 30
        self.batch_size = 16


        #关系抽取配置
        self.bert_dim = 768
        self.num_rel = 18
        self.bert_best_path = os.path.join(current_app.root_path,'bert_pretrain')
        self.tokenizer = BertTokenizer.from_pretrained(self.bert_best_path)
        self.bert_config = BertConfig.from_pretrained(f'{self.bert_best_path}\\bert_config.json')
        self.entertainment_data_path  =  os.path.join(current_app.root_path,'data','entertainment')
        self.bert_train_path = os.path.join(self.entertainment_data_path,'train.json')
        self.bert_dev_path = os.path.join(self.entertainment_data_path, 'dev.json')
        self.bert_test_path = os.path.join(self.entertainment_data_path,'test.json')
        self.bert_rel_path = os.path.join(self.entertainment_data_path,'relation.json')
        #关系类别相关：映射字典，类别数
        self.id2rel = json.load(open(self.bert_rel_path,encoding = 'utf-8'))
        self.rel2id = {v:int(k) for k,v in self.id2rel.items()}
        self.bert_epochs = 10
        self.bert_batch_size = 16
        self.bert_learning_rate = 1e-5
        self.rel_size = len(self.rel2id)